#pragma once
#include <windows.h>
#include <string>
#include <unordered_map>
#include <functional>
#include <stdexcept>
#include <memory>

class DllFunctionBase {
public:
    virtual ~DllFunctionBase() = default;
};

template<typename Func>
class DllFunction : public DllFunctionBase {
public:
    explicit DllFunction(Func func) : m_func(func) {}
    Func get() const { return m_func; }
private:
    Func m_func;
};

class DllWrapper {
public:
    explicit DllWrapper(const std::wstring& dllPath) : m_hDll(nullptr) {
        load(dllPath);
    }

    ~DllWrapper() {
        unload();
    }

    DllWrapper(const DllWrapper&) = delete;
    DllWrapper& operator=(const DllWrapper&) = delete;

    DllWrapper(DllWrapper&& other) noexcept : m_hDll(other.m_hDll), m_functions(std::move(other.m_functions)) {
        other.m_hDll = nullptr;
    }

    DllWrapper& operator=(DllWrapper&& other) noexcept {
        if (this != &other) {
            unload();
            m_hDll = other.m_hDll;
            m_functions = std::move(other.m_functions);
            other.m_hDll = nullptr;
        }
        return *this;
    }

    template<typename Func>
    void registerFunction(const std::string& funcName) {
        if (!m_hDll) {
            throw std::runtime_error("DLL not loaded");
        }

        FARPROC procAddress = GetProcAddress(m_hDll, funcName.c_str());
        if (!procAddress) {
            throw std::runtime_error("Failed to get function address: " + funcName);
        }

        m_functions[funcName] = std::make_unique<DllFunction<Func>>(reinterpret_cast<Func>(procAddress));
    }

    template<typename Func>
    Func getFunction(const std::string& funcName) const {
        auto it = m_functions.find(funcName);
        if (it == m_functions.end()) {
            throw std::runtime_error("Function not registered: " + funcName);
        }

        auto funcPtr = dynamic_cast<DllFunction<Func>*>(it->second.get());
        if (!funcPtr) {
            throw std::runtime_error("Function type mismatch: " + funcName);
        }

        return funcPtr->get();
    }

    bool isLoaded() const {
        return m_hDll != nullptr;
    }

private:
    void load(const std::wstring& dllPath) {
        m_hDll = LoadLibrary(dllPath.c_str());
        if (!m_hDll) {
            throw std::runtime_error("Failed to load DLL: " + std::string(dllPath.begin(), dllPath.end()));
        }
    }

    void unload() {
        if (m_hDll) {
            FreeLibrary(m_hDll);
            m_hDll = nullptr;
        }
    }

    HINSTANCE m_hDll;
    std::unordered_map<std::string, std::unique_ptr<DllFunctionBase>> m_functions;
};    